from qiskit import QuantumCircuit
import numpy as np
from typing import List, Optional


class UniformCircuit:
    """
    A class for creating quantum circuits that produce measurement outcomes
    following a uniform distribution.
    
    This class provides methods to generate quantum circuits whose measurement
    statistics produce a uniform distribution where all possible bit strings
    have equal probability.
    """
    
    def __init__(self, num_qubits: int, qubit_indices: Optional[List[int]] = None):
        """
        Initialize the uniform distribution circuit generator.
        
        Args:
            num_qubits: Total number of qubits in the circuit
            qubit_indices: Specific qubits to use. If None, all qubits are used.
        """
        self.num_qubits = num_qubits
        
        # If qubit_indices is not provided, use all qubits
        if qubit_indices is None:
            self.qubit_indices = list(range(num_qubits))
        else:
            self.qubit_indices = qubit_indices
            
        self.n_distribution_qubits = len(self.qubit_indices)
        self.num_states = 2**self.n_distribution_qubits
        
        self.circuit = None
        
    def generate_circuit(self, measure: bool = True) -> QuantumCircuit:
        """
        Generate a quantum circuit that produces uniform distribution measurement outcomes.
        
        Args:
            measure: Whether to include measurement operations in the circuit
            
        Returns:
            A quantum circuit that will produce uniform distribution outcomes when measured
        """
        # Create a new quantum circuit
        qc = QuantumCircuit(self.num_qubits, len(self.qubit_indices) if measure else 0)
        
        # Apply Hadamard gates to all target qubits to create equal superposition
        # This creates a uniform distribution where each bit string has equal probability
        for idx in self.qubit_indices:
            qc.h(idx)
        
        # Add measurements if requested
        if measure:
            for i, qubit_idx in enumerate(self.qubit_indices):
                qc.measure(qubit_idx, i)
        
        self.circuit = qc
        return qc
    
    def generate_alternative_circuit(self, measure: bool = True) -> QuantumCircuit:
        """
        Generate a quantum circuit that produces uniform distribution using the initialize method.
        This is an alternative approach that's equivalent to using Hadamard gates.
        
        Args:
            measure: Whether to include measurement operations in the circuit
            
        Returns:
            A quantum circuit that will produce uniform distribution outcomes when measured
        """
        # Create a new quantum circuit
        qc = QuantumCircuit(self.num_qubits, len(self.qubit_indices) if measure else 0)
        
        # For uniform distribution, all amplitudes are equal
        # With N qubits, each amplitude is 1/sqrt(2^N)
        amplitudes = np.ones(self.num_states) / np.sqrt(self.num_states)
        
        # Initialize the target qubits with uniform amplitudes
        target_qubits = self.qubit_indices[:self.n_distribution_qubits]
        qc.initialize(amplitudes, target_qubits)
        
        # Add measurements if requested
        if measure:
            for i, qubit_idx in enumerate(self.qubit_indices):
                qc.measure(qubit_idx, i)
        
        self.circuit = qc
        return qc


if __name__ == "__main__":
    # Example usage
    uniform = UniformCircuit(num_qubits=4, qubit_indices=[0, 1, 2, 3])
    
    # Method 1: Using Hadamard gates (standard approach)
    circuit = uniform.generate_circuit()
    print("Uniform distribution circuit using Hadamard gates:")
    print(circuit.draw())
    
    # Method 2: Using initialize method (alternative approach)
    circuit_alt = uniform.generate_alternative_circuit()
    print("\nUniform distribution circuit using initialize method:")
    print(circuit_alt.draw())
    
    # Expected outcome: With 4 qubits, each of the 16 possible bit strings (0000 to 1111)
    # will have equal probability of 1/16 when measured 